package com.rj.framework.compiler;

import javax.tools.*;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.OutputStream;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.URI;
import java.util.*;

/**
 * 动态编译器
 */
public class DynamicCompiler {
    private final JavaCompiler compiler;
    private final StandardJavaFileManager fileManager;
    private final List<File> outputDirs;

    public DynamicCompiler() {
        this.compiler = ToolProvider.getSystemJavaCompiler();
        this.fileManager = compiler.getStandardFileManager(null, null, null);
        this.outputDirs = new ArrayList<>();
    }

    public Class<?> compile(String className, String sourceCode) {
        try {
            // 创建内存中的Java源文件
            JavaSourceFromString src = new JavaSourceFromString(className, sourceCode);
            // 设置编译选项
            List<String> options = new ArrayList<>();
            options.add("-classpath");
            options.add(System.getProperty("java.class.path"));
            // 设置输出目录（内存）
            MemoryJavaFileManager memFileManager = new MemoryJavaFileManager(fileManager);
            // 执行编译
            JavaCompiler.CompilationTask task = compiler.getTask(
                    null,
                    memFileManager,
                    null,
                    options,
                    null,
                    Collections.singletonList(src)
            );
            if (!task.call()) {
                throw new RuntimeException("Compilation failed");
            }
            // 加载编译后的类
            MemoryClassLoader classLoader = new MemoryClassLoader(memFileManager);
            return classLoader.loadClass(className);
        } catch (Exception e) {
            throw new RuntimeException("Compilation error", e);
        }
    }

    static class JavaSourceFromString extends SimpleJavaFileObject {
        final String code;

        JavaSourceFromString(String name, String code) {
            super(URI.create("string:///" + name.replace('.', '/') + Kind.SOURCE.extension),
                    Kind.SOURCE);
            this.code = code;
        }

        @Override
        public CharSequence getCharContent(boolean ignoreEncodingErrors) {
            return code;
        }
    }

    static class MemoryJavaFileManager extends ForwardingJavaFileManager<StandardJavaFileManager> {
        private final Map<String, MemoryJavaFileObject> classFiles = new HashMap<>();

        MemoryJavaFileManager(StandardJavaFileManager fileManager) {
            super(fileManager);
        }

        @Override
        public JavaFileObject getJavaFileForOutput(Location location, String className,
                                                   JavaFileObject.Kind kind, FileObject sibling) throws IOException {
            if (kind == JavaFileObject.Kind.CLASS) {
                MemoryJavaFileObject file = new MemoryJavaFileObject(className);
                classFiles.put(className, file);
                return file;
            }
            return super.getJavaFileForOutput(location, className, kind, sibling);
        }

        public Map<String, MemoryJavaFileObject> getClassFiles() {
            return classFiles;
        }
    }

    static class MemoryJavaFileObject extends SimpleJavaFileObject {
        private final ByteArrayOutputStream bos = new ByteArrayOutputStream();
        private final String className;

        MemoryJavaFileObject(String className) {
            super(URI.create("byte:///" + className.replace('.', '/') + Kind.CLASS.extension),
                    Kind.CLASS);
            this.className = className;
        }

        @Override
        public OutputStream openOutputStream() {
            return bos;
        }

        public byte[] getBytes() {
            return bos.toByteArray();
        }

        public String getClassName() {
            return className;
        }
    }

    static class MemoryClassLoader extends ClassLoader {
        private final Map<String, MemoryJavaFileObject> classFiles;

        MemoryClassLoader(MemoryJavaFileManager fileManager) {
            this.classFiles = fileManager.getClassFiles();
        }

        @Override
        protected Class<?> findClass(String name) throws ClassNotFoundException {
            MemoryJavaFileObject file = classFiles.get(name);
            if (file != null) {
                byte[] bytes = file.getBytes();
                return defineClass(name, bytes, 0, bytes.length);
            }
            return super.findClass(name);
        }
    }

    public static void main(String[] args) throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
        DynamicCompiler compiler = new DynamicCompiler();
        System.out.println("");
        Class<?> clazz = compiler.compile("com.Test",
                "package com;\n" +
                        "public class Test{" +
                        "public static void run() {" +
                        "System.out.println(\"Hello World!\");" +
                        "}}");
        Method run = clazz.getMethod("run");
        run.invoke(null);
    }
}